import math
import tkinter
# import matplotlib
# matplotlib.use('TkAgg')
import os
import warnings

import CR_DSPPytorch
from ERPNLC_TDC import simulation_ERP_NLC_module

os.environ["OMP_NUM_THREADS"] = "8"
from inspect import currentframe, getframeinfo
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import scipy.io as scio
from CR_DSPPytorch import EDCLayer, PhaseRecLayer, FIRLayer, format_rt, cmul, OFDE
import util

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
BASE_DIR = os.path.dirname(__file__)


def OverlapCDC(lp_range=[1], fft_size_list=[2048], save_result_flag=False):
    resultPath = os.path.join(BASE_DIR, 'result/FDE')
    os.makedirs(resultPath, exist_ok=True)
    pick_syms_num = 131072
    modOrder = 4
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    BERcache = np.zeros([len(lp_range), len(fft_size_list)])
    Qcache = np.zeros_like(BERcache)
    Complexitycache = np.zeros_like(BERcache)
    save_result_flag = save_result_flag

    for lpIndx, lp in enumerate(lp_range):
        for tapIndx, tap in enumerate(fft_size_list):
            # print(f'lp={lp}, fft_size={tap}')
            # load data
            dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                              f'/trSet_lp_{lp}.mat')
            data = scio.loadmat(dataPath)
            symbolRate = data['symbolRate'][0, 0]
            spanNum = data['spanNum'][0, 0]
            spanLen = data['spanLen'][0, 0]
            sig = data['sig']
            prbs = data['prbs']
            L = spanNum * spanLen
            D = 17e-6
            DL = D * L / 2

            minimum_overlap = calculate_minimum_overlap(symbolRate, DL)
            print(f'minimum_overlap for OFDE: {minimum_overlap}')
            '''Calculate the complexity for FDE'''
            Comp = tap * (6 * 1 / 2 * math.log2(tap) + 3) / (tap * 1 / 2 + 1)

            sig = torch.from_numpy(sig[np.newaxis, ...])
            constellations = torch.from_numpy(util.CONST_16QAM)

            '''DSP procedure'''
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)
            sig = sig.to(chosen_device)
            lms = lms.to(chosen_device)
            sig = sig[..., 0:pick_syms_num * 2]
            prbs = prbs[..., 0:pick_syms_num * modOrder]

            sig = torch.view_as_real(sig)
            sig = OFDE(sig, symbolRate, Ntap=tap, DL=DL)
            sig = pr(sig[..., 1::2, :])
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            sig = torch.view_as_complex(sig)
            sig = sig.cpu().data.numpy().squeeze()
            # ber, Q = simulation_ERP_NLC_module(sig, prbs, lp, symbolRate, spanLen, spanNum)
            sig, ber, _ = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())
            ber = np.mean(ber)
            Q = util.ber2q(ber)
            BERcache[lpIndx, tapIndx] = ber
            Qcache[lpIndx, tapIndx] = Q
            Complexitycache[lpIndx, tapIndx] = Comp
            print(f'lp: {lp}, fft_size:{tap}, BER: {ber}, Q: {Q}, Complexity: {Comp}')

    if save_result_flag:
        with open(os.path.join(resultPath, f'FDE_CDC_Complexity_Vary_lp.csv'), 'w') as f:
            f.write('lp,fft_size,Complexity,BER,Q\n')
            for lpIndx, lp in enumerate(lp_range):
                for tapIndx, tap in enumerate(fft_size_list):
                    f.write(
                        f'{lp},{tap},{Complexitycache[lpIndx, tapIndx]},{BERcache[lpIndx, tapIndx]},{Qcache[lpIndx, tapIndx]}\n')
        print(f'result saved in \"{resultPath}\"')


def calculate_minimum_overlap(symbol_rate, DL, wave_length=1552.52e-9, sample_factor=2):
    c = 3e8
    Ts = 1 / symbol_rate / sample_factor
    k = np.pi * c * Ts ** 2
    m = wave_length ** 2 * DL
    N_overlap = 2 * math.floor(1 / k * np.sqrt(k ** 2 + 4 * m ** 2)) + 2
    return N_overlap


if __name__ == '__main__':
    OverlapCDC(lp_range=[-2, -1, 0, 1, 2, 3],
               fft_size_list=[2048],
               save_result_flag=True)
